# -*- coding: utf-8 -*-
"""
遗传算法脚本

使用NSGA-III多目标遗传算法寻找最佳参数组合，以优化模型性能。
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import mean_squared_error, r2_score
from model import create_model, load_model
import pickle
import json
import time
from pymoo.core.problem import Problem
from pymoo.algorithms.moo.nsga3 import NSGA3
from pymoo.util.ref_dirs import get_reference_directions
from pymoo.optimize import minimize
from pymoo.core.callback import Callback
from pymoo.visualization.scatter import Scatter
from pymoo.indicators.hv import HV
import torch

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 创建输出目录
if not os.path.exists('output'):
    os.makedirs('output')

# 创建模型目录
if not os.path.exists('models'):
    os.makedirs('models')

# 创建图表目录
if not os.path.exists('figures'):
    os.makedirs('figures')

# 创建遗传算法结果目录
if not os.path.exists('ga_results'):
    os.makedirs('ga_results')


class MaterialOptimizationProblem(Problem):
    """
    材料优化问题
    
    使用NSGA-III多目标遗传算法寻找最佳参数组合，以优化太阳反射率和红外发射率。
    
    参数:
        model (nn.Module): 训练好的神经网络模型
        feature_names (list): 特征名称列表
        encoders (dict): 特征编码器字典
        n_var (int): 变量数量
        n_obj (int): 目标数量
        xl (np.ndarray): 变量下界
        xu (np.ndarray): 变量上界
        var_names (list): 变量名称列表
        var_types (list): 变量类型列表
    """
    def __init__(self, model, feature_names, encoders, n_var, n_obj, xl, xu, var_names, var_types):
        super().__init__(n_var=n_var, n_obj=n_obj, xl=xl, xu=xu)
        
        self.model = model
        self.feature_names = feature_names
        self.encoders = encoders
        self.var_names = var_names
        self.var_types = var_types  # 'categorical' 或 'continuous'
        
        # 定义必须有值的输入列和可以为空的输入列
        self.required_input_cols = ['高分子基体（缩写）', '无机添加剂（缩写）']
        self.optional_input_cols = ['其余材料', '加工方法']
        
        # 将模型设置为评估模式
        self.model.eval()
    
    def _evaluate(self, x, out, *args, **kwargs):
        """
        评估函数
        
        参数:
            x (np.ndarray): 输入参数
            out (dict): 输出结果
        """
        # 将参数转换为特征向量
        X = np.zeros((x.shape[0], len(self.feature_names)))
        
        # 处理每个样本
        for i in range(x.shape[0]):
            # 构建输入特征
            features = self._params_to_features(x[i])
            X[i] = features
        
        # 转换为PyTorch张量
        X_tensor = torch.FloatTensor(X)
        
        # 使用模型预测
        with torch.no_grad():
            predictions = self.model(X_tensor).numpy()
        
        # 目标是最大化太阳反射率和红外发射率，所以取负值（因为pymoo是最小化问题）
        f1 = -predictions[:, 0]  # 太阳反射率（取负值）
        f2 = -predictions[:, 1]  # 红外发射率（取负值）
        
        # 组合目标
        out["F"] = np.column_stack([f1, f2])
    
    def _params_to_features(self, params):
        """
        将参数转换为特征向量
        
        参数:
            params (np.ndarray): 参数数组
            
        返回:
            np.ndarray: 特征向量
        """
        # 初始化特征向量
        features = np.zeros(len(self.feature_names))
        
        # 处理每个参数
        param_idx = 0
        for col, encoder in self.encoders.items():
            # 获取该特征的类别数量
            n_categories = len(encoder.categories_[0])
            
            # 检查变量类型
            var_type = self.var_types[param_idx]
            
            if var_type == 'categorical':
                # 对于分类变量，选择一个类别
                category_idx = int(params[param_idx] * n_categories)
                if category_idx >= n_categories:
                    category_idx = n_categories - 1
                
                # 设置对应的one-hot编码
                start_idx = 0
                for prev_col, prev_encoder in self.encoders.items():
                    if prev_col == col:
                        break
                    start_idx += len(prev_encoder.categories_[0])
                
                # 对于可选输入列，我们需要特殊处理
                if col in self.optional_input_cols:
                    # 检查是否选择了'缺失'类别
                    categories = encoder.categories_[0]
                    selected_category = categories[category_idx]
                    
                    # 无论是否选择了'缺失'类别，都将其作为有效特征值处理
                    # 设置选中的类别为1，包括'缺失'类别
                    features[start_idx + category_idx] = 1
                else:
                    # 对于必须有值的输入列，直接设置选中的类别为1
                    features[start_idx + category_idx] = 1
            else:  # 'continuous'
                # 对于连续变量，直接使用参数值
                # 这里假设连续变量已经被归一化到[0, 1]范围
                features[param_idx] = params[param_idx]
            
            param_idx += 1
        
        return features


class GACallback(Callback):
    """
    遗传算法回调函数
    
    用于记录和可视化遗传算法的进度。
    """
    def __init__(self, problem, save_dir='ga_results'):
        super().__init__()
        self.problem = problem
        self.save_dir = save_dir
        self.data = []
        
        # 确保保存目录存在
        os.makedirs(save_dir, exist_ok=True)
    
    def notify(self, algorithm):
        # 获取当前代的非支配解
        pop = algorithm.pop
        F = pop.get("F")
        X = pop.get("X")
        
        # 记录数据
        self.data.append({
            'gen': algorithm.n_gen,
            'F': F.copy(),
            'X': X.copy()
        })
        
        # 每10代保存一次结果
        if algorithm.n_gen % 10 == 0:
            # 可视化当前代的非支配解
            self._visualize_generation(algorithm.n_gen, F, X)
            
            # 保存当前代的结果
            self._save_generation(algorithm.n_gen, F, X)
    
    def _visualize_generation(self, gen, F, X):
        """
        可视化当前代的非支配解
        """
        # 创建散点图
        plt.figure(figsize=(10, 8))
        plt.scatter(-F[:, 0], -F[:, 1], s=30)
        plt.xlabel('太阳反射率')
        plt.ylabel('红外发射率')
        plt.title(f'第 {gen} 代非支配解')
        plt.grid(True)
        plt.savefig(f"{self.save_dir}/gen_{gen}_pareto.png", dpi=300, bbox_inches='tight')
        plt.close()
    
    def _save_generation(self, gen, F, X):
        """
        保存当前代的结果
        """
        # 创建结果DataFrame
        results = []
        for i in range(len(X)):
            # 将参数转换为有意义的值
            params = {}
            for j, name in enumerate(self.problem.var_names):
                params[name] = X[i, j]
            
            # 添加目标值
            params['太阳反射率'] = -F[i, 0]
            params['红外发射率'] = -F[i, 1]
            
            results.append(params)
        
        # 创建DataFrame并保存
        df = pd.DataFrame(results)
        df.to_csv(f"{self.save_dir}/gen_{gen}_results.csv", index=False, encoding='utf-8-sig')


def load_data_and_model():
    """
    加载数据和模型
    
    返回:
        tuple: (model, feature_names, encoders)
    """
    print("加载数据和模型...")
    
    # 加载特征名称
    with open('output/feature_names.pkl', 'rb') as f:
        feature_names = pickle.load(f)
    
    # 加载编码器
    try:
        with open('output/encoders.pkl', 'rb') as f:
            encoders = pickle.load(f)
    except Exception as e:
        print(f"加载编码器时出错: {e}")
        # 尝试处理可能的兼容性问题
        if isinstance(e, ModuleNotFoundError) and "numpy._core" in str(e):
            print("尝试解决numpy版本兼容性问题...")
            import sys
            import numpy
            sys.modules['numpy._core'] = numpy
            with open('output/encoders.pkl', 'rb') as f:
                encoders = pickle.load(f)
    
    # 加载最佳模型信息
    with open('output/best_model_info.json', 'r') as f:
        best_model_info = json.load(f)
    
    best_model_name = best_model_info['name']
    print(f"最佳模型: {best_model_name}")
    
    # 加载模型配置
    model_configs = {
        'mlp_model': {
            'type': 'mlp',
            'params': {
                'hidden_dims': [64, 32],
                'dropout_rate': 0.2,
                'use_batch_norm': True,
                'activation': 'relu'
            }
        },
        'residual_model': {
            'type': 'residual',
            'params': {
                'hidden_dims': [64, 128, 64],
                'num_res_blocks': 2,
                'dropout_rate': 0.2,
                'use_batch_norm': True
            }
        },
        'dual_branch_model': {
            'type': 'dual_branch',
            'params': {
                'hidden_dims': [64, 128],
                'branch_dims': [64, 32],
                'dropout_rate': 0.2,
                'use_batch_norm': True,
                'num_heads': 4
            }
        },
        'ensemble_model': {
            'type': 'ensemble',
            'params': {
                'num_models': 3,
                'model_type': 'dual_branch',
                'hidden_dims': [64, 128],
                'dropout_rate': 0.2,
                'use_batch_norm': True
            }
        }
    }
    
    # 获取模型配置
    model_config = model_configs[best_model_name]
    
    # 加载模型
    input_dim = len(feature_names)
    model = load_model(
        model_config['type'],
        input_dim,
        f"models/{best_model_name}.pth",
        **model_config['params']
    )
    
    # 设置为评估模式
    model.eval()
    
    return model, feature_names, encoders


def define_optimization_problem(model, feature_names, encoders):
    """
    定义优化问题
    
    参数:
        model (nn.Module): 训练好的神经网络模型
        feature_names (list): 特征名称列表
        encoders (dict): 特征编码器字典
        
    返回:
        MaterialOptimizationProblem: 优化问题实例
    """
    print("定义优化问题...")
    
    # 定义变量（参数）
    var_names = list(encoders.keys())  # 使用编码器的键作为变量名
    n_var = len(var_names)
    
    # 定义必须有值的输入列和可以为空的输入列
    required_input_cols = ['高分子基体（缩写）', '无机添加剂（缩写）']
    optional_input_cols = ['其余材料', '加工方法']
    
    # 定义变量类型（这里假设所有变量都是分类变量）
    var_types = ['categorical'] * n_var
    
    # 定义变量范围（对于分类变量，范围是[0, 1)，将在内部转换为类别索引）
    xl = np.zeros(n_var)
    xu = np.ones(n_var)
    
    # 定义目标数量（太阳反射率和红外发射率）
    n_obj = 2
    
    # 创建优化问题
    problem = MaterialOptimizationProblem(
        model, feature_names, encoders, n_var, n_obj, xl, xu, var_names, var_types
    )
    
    # 打印编码器信息，帮助理解特征编码
    print("\n编码器信息:")
    for col, encoder in encoders.items():
        categories = encoder.categories_[0]
        print(f"  {col}: {len(categories)} 个类别")
        if col in optional_input_cols:
            # 检查是否有'缺失'类别
            if '缺失' in categories:
                print(f"    包含'缺失'类别，可以生成缺失值")
            else:
                print(f"    警告: '{col}' 没有'缺失'类别，可能无法正确处理缺失值")
    
    return problem


def run_genetic_algorithm(problem, n_gen=100, pop_size=100):
    """
    运行遗传算法
    
    参数:
        problem (Problem): 优化问题
        n_gen (int): 迭代代数
        pop_size (int): 种群大小
        
    返回:
        tuple: (res, callback)
    """
    print(f"运行遗传算法，代数: {n_gen}，种群大小: {pop_size}...")
    
    # 获取参考方向（用于NSGA-III）
    ref_dirs = get_reference_directions("das-dennis", problem.n_obj, n_partitions=12)
    
    # 创建NSGA-III算法
    algorithm = NSGA3(
        pop_size=pop_size,
        ref_dirs=ref_dirs
    )
    
    # 创建回调
    callback = GACallback(problem)
    
    # 运行优化
    start_time = time.time()
    res = minimize(
        problem,
        algorithm,
        ('n_gen', n_gen),
        seed=42,
        save_history=True,
        verbose=False,
        callback=callback
    )
    
    # 计算运行时间
    run_time = time.time() - start_time
    print(f"遗传算法运行完成，耗时: {run_time:.2f} 秒")
    
    return res, callback


def analyze_results(res, problem, callback):
    """
    分析遗传算法结果
    
    参数:
        res (Result): 优化结果
        problem (Problem): 优化问题
        callback (GACallback): 回调实例
    """
    print("分析遗传算法结果...")
    
    # 获取最终的非支配解
    X = res.X
    F = res.F
    
    # 转换为实际的目标值（取负值，因为我们是最大化问题）
    F_actual = -F
    
    # 创建结果DataFrame
    results = []
    for i in range(len(X)):
        # 将参数转换为有意义的值
        params = {}
        for j, name in enumerate(problem.var_names):
            # 对于分类变量，转换为类别名称
            if problem.var_types[j] == 'categorical':
                encoder = problem.encoders[name]
                categories = encoder.categories_[0]
                category_idx = int(X[i, j] * len(categories))
                if category_idx >= len(categories):
                    category_idx = len(categories) - 1
                params[name] = categories[category_idx]
            else:  # 'continuous'
                params[name] = X[i, j]
        
        # 添加目标值
        params['太阳反射率'] = F_actual[i, 0]
        params['红外发射率'] = F_actual[i, 1]
        params['综合性能'] = (F_actual[i, 0] + F_actual[i, 1]) / 2  # 简单的综合性能指标
        
        results.append(params)
    
    # 创建DataFrame
    df = pd.DataFrame(results)
    
    # 按综合性能排序
    df = df.sort_values('综合性能', ascending=False)
    
    # 保存结果
    df.to_csv('ga_results/final_results.csv', index=False, encoding='utf-8-sig')
    print(f"最终结果已保存到 ga_results/final_results.csv")
    
    # 可视化最终的非支配解
    plt.figure(figsize=(10, 8))
    plt.scatter(F_actual[:, 0], F_actual[:, 1], s=30)
    plt.xlabel('太阳反射率')
    plt.ylabel('红外发射率')
    plt.title('最终非支配解')
    plt.grid(True)
    plt.savefig("ga_results/final_pareto.png", dpi=300, bbox_inches='tight')
    
    # 可视化优化过程
    visualize_optimization_process(callback)
    
    # 分析最佳解
    analyze_best_solutions(df)


def visualize_optimization_process(callback):
    """
    可视化优化过程
    
    参数:
        callback (GACallback): 回调实例
    """
    print("可视化优化过程...")
    
    # 提取每代的数据
    generations = [data['gen'] for data in callback.data]
    
    # 计算每代的超体积指标（如果可用）
    try:
        # 创建参考点（比最差的解还要差一点）
        # 注意：由于我们是最大化问题，但pymoo默认是最小化问题，所以F值是负的
        # 因此参考点应该是负的最小值，即比最差的解还要差
        worst_f = np.max([data['F'] for data in callback.data], axis=0)
        ref_point = worst_f * 1.1  # 比最差的解还要差10%
        
        # 计算每代的超体积
        hv = []
        for data in callback.data:
            F = data['F']  # 使用原始F值（负值）
            indicator = HV(ref_point=ref_point)
            hv.append(indicator(F))
           
        
        # 绘制超体积随代数的变化
        plt.figure(figsize=(10, 6))
        plt.plot(generations, hv, marker='o')
        plt.xlabel('代数')
        plt.ylabel('超体积指标')
        plt.title('超体积指标随代数的变化')
        plt.grid(True)
        plt.savefig("ga_results/hypervolume.png", dpi=300, bbox_inches='tight')
    except Exception as e:
        print(f"计算超体积指标时出错: {e}")
    
    # 可视化每代的非支配解分布变化
    try:
        plt.figure(figsize=(12, 10))
        
        # 选择几个关键代数进行可视化
        key_gens = [0, len(generations)//4, len(generations)//2, len(generations)-1]
        
        for i, gen_idx in enumerate(key_gens):
            if gen_idx < len(callback.data):
                data = callback.data[gen_idx]
                F = -data['F']  # 转换回正值
                
                plt.subplot(2, 2, i+1)
                plt.scatter(F[:, 0], F[:, 1], s=30)
                plt.xlabel('太阳反射率')
                plt.ylabel('红外发射率')
                plt.title(f'第 {data["gen"]} 代非支配解')
                plt.grid(True)
        
        plt.tight_layout()
        plt.savefig("ga_results/pareto_evolution.png", dpi=300, bbox_inches='tight')
    except Exception as e:
        print(f"可视化非支配解分布时出错: {e}")


def analyze_best_solutions(df):
    """
    分析最佳解
    
    参数:
        df (pd.DataFrame): 结果DataFrame
    """
    print("分析最佳解...")
    
    # 获取前10个最佳解
    top_solutions = df.head(10)
    
    # 保存前10个最佳解
    top_solutions.to_csv('ga_results/top_solutions.csv', index=False, encoding='utf-8-sig')
    print(f"前10个最佳解已保存到 ga_results/top_solutions.csv")
    
    # 分析参数分布
    param_cols = [col for col in df.columns if col not in ['太阳反射率', '红外发射率', '综合性能']]
    
    # 对于每个参数，分析其在前10个最佳解中的分布
    param_analysis = {}
    for col in param_cols:
        # 获取参数在前10个最佳解中的值
        values = top_solutions[col].value_counts().to_dict()
        param_analysis[col] = values
    
    # 保存参数分析结果
    with open('ga_results/param_analysis.json', 'w', encoding='utf-8') as f:
        json.dump(param_analysis, f, ensure_ascii=False, indent=4)
    
    # 可视化参数分布
    try:
        # 对于每个参数，创建一个条形图
        for col in param_cols:
            plt.figure(figsize=(10, 6))
            values = top_solutions[col].value_counts()
            values.plot(kind='bar')
            plt.xlabel(col)
            plt.ylabel('频数')
            plt.title(f'{col} 在前10个最佳解中的分布')
            plt.tight_layout()
            plt.savefig(f"ga_results/{col}_distribution.png", dpi=300, bbox_inches='tight')
            plt.close()
    except Exception as e:
        print(f"可视化参数分布时出错: {e}")
    
    # 分析最佳解的特点
    print("\n最佳解的特点:")
    best_solution = df.iloc[0]
    print(f"太阳反射率: {best_solution['太阳反射率']:.4f}")
    print(f"红外发射率: {best_solution['红外发射率']:.4f}")
    print(f"综合性能: {best_solution['综合性能']:.4f}")
    
    print("\n最佳参数组合:")
    for col in param_cols:
        print(f"{col}: {best_solution[col]}")


def main():
    # 加载数据和模型
    model, feature_names, encoders = load_data_and_model()
    
    # 定义优化问题
    problem = define_optimization_problem(model, feature_names, encoders)
    
    # 运行遗传算法
    res, callback = run_genetic_algorithm(problem, n_gen=100, pop_size=100)
    
    # 分析结果
    analyze_results(res, problem, callback)
    
    print("\n遗传算法优化完成！")
    print(f"结果已保存到 ga_results/ 目录")


if __name__ == "__main__":
    main()